import pandas as pd
import numpy as np
import yfinance as yf  # 用于获取股票数据
import akshare as ak  # 用于获取中国期货数据
import talib as ta  # 技术指标计算库
from datetime import datetime, timedelta

def get_stock_data(symbol, period="1y", interval="1d"):
    """
    获取股票K线数据
    :param symbol: 股票代码 (格式: 美股AAPL, A股600000.SS)
    :param period: 数据周期 (1d,5d,1mo,3mo,6mo,1y,2y,5y,10y,ytd,max)
    :param interval: K线间隔 (1m,2m,5m,15m,30m,60m,90m,1h,1d,5d,1wk,1mo,3mo)
    :return: 包含K线数据和技术指标的DataFrame
    """
    try:
        # 从Yahoo Finance下载数据
        stock = yf.Ticker(symbol)
        df = stock.history(period=period, interval=interval)
        
        # 重置索引并重命名列
        df.reset_index(inplace=True)
        df.rename(columns={
            'Date': 'timestamp',
            'Open': 'open',
            'High': 'high',
            'Low': 'low',
            'Close': 'close',
            'Volume': 'volume'
        }, inplace=True)
        
        # 添加股票代码列
        df['symbol'] = symbol
        
        return df
    except Exception as e:
        print(f"获取股票数据失败: {e}")
        return pd.DataFrame()

def get_futures_data(symbol, start_date, end_date, period="daily"):
    """
    获取期货K线数据
    :param symbol: 期货合约代码 (格式: RB9999 螺纹钢主力合约)
    :param start_date: 开始日期 (YYYY-MM-DD)
    :param end_date: 结束日期 (YYYY-MM-DD)
    :param period: 数据周期 (daily, weekly, monthly)
    :return: 包含K线数据和技术指标的DataFrame
    """
    try:
        # 使用AKShare获取期货数据
        df = ak.futures_zh_spot(symbol=symbol, period=period, 
                               start_date=start_date, end_date=end_date)
        
        # 重命名列
        df.rename(columns={
            '日期': 'timestamp',
            '开盘': 'open',
            '最高': 'high',
            '最低': 'low',
            '收盘': 'close',
            '成交量': 'volume'
        }, inplace=True)
        
        # 添加期货代码列
        df['symbol'] = symbol
        
        # 转换日期格式
        df['timestamp'] = pd.to_datetime(df['timestamp'])
        
        # 转换数据类型
        df['open'] = df['open'].astype(float)
        df['high'] = df['high'].astype(float)
        df['low'] = df['low'].astype(float)
        df['close'] = df['close'].astype(float)
        df['volume'] = df['volume'].astype(float)
        
        return df
    except Exception as e:
        print(f"获取期货数据失败: {e}")
        return pd.DataFrame()

def calculate_technical_indicators(df):
    """
    计算技术指标并添加到DataFrame
    :param df: 包含K线数据的DataFrame
    :return: 添加技术指标后的DataFrame
    """
    if df.empty:
        return df
    
    # 计算移动平均线
    df['MA5'] = ta.SMA(df['close'], timeperiod=5)
    df['MA20'] = ta.SMA(df['close'], timeperiod=20)
    df['MA60'] = ta.SMA(df['close'], timeperiod=60)
    
    # 计算相对强弱指数(RSI)
    df['RSI14'] = ta.RSI(df['close'], timeperiod=14)
    
    # 计算MACD
    macd, macd_signal, macd_hist = ta.MACD(df['close'], 
                                          fastperiod=12, 
                                          slowperiod=26, 
                                          signalperiod=9)
    df['MACD'] = macd
    df['MACD_Signal'] = macd_signal
    df['MACD_Hist'] = macd_hist
    
    # 计算布林带
    upper, middle, lower = ta.BBANDS(df['close'], 
                                    timeperiod=20, 
                                    nbdevup=2, 
                                    nbdevdn=2, 
                                    matype=0)
    df['BOLL_Upper'] = upper
    df['BOLL_Middle'] = middle
    df['BOLL_Lower'] = lower
    
    # 计算随机指标(KDJ)
    slowk, slowd = ta.STOCH(df['high'], df['low'], df['close'],
                           fastk_period=9, 
                           slowk_period=3, 
                           slowk_matype=0, 
                           slowd_period=3, 
                           slowd_matype=0)
    df['KDJ_K'] = slowk
    df['KDJ_D'] = slowd
    df['KDJ_J'] = 3 * slowk - 2 * slowd
    
    return df

def save_to_csv(df, filename):
    """
    保存数据到CSV文件
    :param df: 包含数据的DataFrame
    :param filename: 保存的文件名
    """
    try:
        # 确保数据按时间排序
        df.sort_values('timestamp', inplace=True)
        
        # 保存到CSV
        df.to_csv(filename, index=False, encoding='utf-8-sig')
        print(f"数据已保存至: {filename}")
    except Exception as e:
        print(f"保存数据失败: {e}")

def main():
    # 设置日期范围
    end_date = datetime.now().strftime('%Y-%m-%d')
    start_date = (datetime.now() - timedelta(days=365)).strftime('%Y-%m-%d')
    
    # 获取股票数据 (示例: 苹果股票)
    stock_df = get_stock_data("AAPL", period="1y", interval="1d")
    if not stock_df.empty:
        stock_df = calculate_technical_indicators(stock_df)
        save_to_csv(stock_df, "AAPL_stock_data.csv")
    
    # 获取期货数据 (示例: 螺纹钢主力合约)
    futures_df = get_futures_data("RB9999", start_date, end_date, "daily")
    if not futures_df.empty:
        futures_df = calculate_technical_indicators(futures_df)
        save_to_csv(futures_df, "RB9999_futures_data.csv")
    
    # 组合数据 (如果需要组合多个资产)
    if not stock_df.empty and not futures_df.empty:
        combined_df = pd.concat([stock_df, futures_df], ignore_index=True)
        save_to_csv(combined_df, "combined_data.csv")

if __name__ == "__main__":
    main()